{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7ab0c751",
   "metadata": {},
   "source": [
    "# {func}`tvm.compile`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bbde515d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import relax, te\n",
    "from tvm.runtime import Executable\n",
    "from tvm.script import ir as I\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2fc48a3",
   "metadata": {},
   "source": [
    "## 测试 TIR 输入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "44ff37b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = te.var(\"n\")\n",
    "A = te.placeholder((n,), name=\"A\")\n",
    "B = te.placeholder((n,), name=\"B\")\n",
    "C = te.compute(A.shape, lambda i: A[i] + B[i], name=\"C\")\n",
    "func = te.create_prim_func([A, B, C])\n",
    "\n",
    "# Test compile with PrimFunc\n",
    "exec_prim = tvm.compile(func)\n",
    "assert isinstance(exec_prim, Executable)\n",
    "\n",
    "# Test compile with IRModule containing PrimFunc\n",
    "mod = tvm.IRModule.from_expr(func)\n",
    "exec_mod = tvm.compile(mod)\n",
    "assert isinstance(exec_mod, Executable)\n",
    "\n",
    "# Verify the compiled module works\n",
    "dev = tvm.cpu(0)\n",
    "a_np = np.random.uniform(size=10).astype(np.float32)\n",
    "b_np = np.random.uniform(size=10).astype(np.float32)\n",
    "a = tvm.nd.array(a_np, dev)\n",
    "b = tvm.nd.array(b_np, dev)\n",
    "c = tvm.nd.array(np.zeros(10, dtype=np.float32), dev)\n",
    "\n",
    "exec_prim(a, b, c)\n",
    "np.testing.assert_allclose(c.numpy(), a_np + b_np)\n",
    "exec_mod(a, b, c)\n",
    "np.testing.assert_allclose(c.numpy(), a_np + b_np)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02e57503",
   "metadata": {},
   "source": [
    "## 测试 Relax 输入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a86631bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a simple Relax program\n",
    "@I.ir_module\n",
    "class MyModule:\n",
    "    @R.function\n",
    "    def main(x: R.Tensor((3, 4), \"float32\"), y: R.Tensor((3, 4), \"float32\")) -> R.Tensor:\n",
    "        z = R.add(x, y)\n",
    "        return z\n",
    "\n",
    "# Test compile with Relax module\n",
    "target = tvm.target.Target(\"llvm\")\n",
    "exec_relax = tvm.compile(MyModule, target)\n",
    "assert isinstance(exec_relax, Executable)\n",
    "\n",
    "# Verify the compiled module works\n",
    "dev = tvm.cpu(0)\n",
    "x_np = np.random.uniform(size=(3, 4)).astype(np.float32)\n",
    "y_np = np.random.uniform(size=(3, 4)).astype(np.float32)\n",
    "x = tvm.nd.array(x_np, dev)\n",
    "y = tvm.nd.array(y_np, dev)\n",
    "\n",
    "vm = tvm.runtime.vm.VirtualMachine(exec_relax, dev)\n",
    "z = vm[\"main\"](x, y)\n",
    "np.testing.assert_allclose(z.numpy(), x_np + y_np)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff5e5f2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
